import os, json
import argparse
import matplotlib.pyplot as plt



# -- coding: utf-8 --
'''
# @Time : 2024/11/24 15:29
# @Author : Shiyu He
# @University : Xinjiang University
'''

parser = argparse.ArgumentParser(description="plot dynamic speed control")
parser.add_argument("input_path", type=str, help="Path to the input json file.")
args = parser.parse_args()


# Function to load data from JSON file
def load_data_from_json(file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

# Load speed and control data from JSON files (replace 'speed.json' and 'control.json' with actual file paths)
speed_data = load_data_from_json(os.path.join(args.input_path, "state.json"))  # Replace with your speed data JSON file path
control_data = load_data_from_json(os.path.join(args.input_path, "control.json"))  # Replace with your control data JSON file path

# Extract speed data
timestamps = speed_data["plane_speeds"]
x_speeds = [speed_data["plane_speeds"][t][0] for t in timestamps]
y_speeds = [speed_data["plane_speeds"][t][1] for t in timestamps]
z_speeds = [speed_data["plane_speeds"][t][2] for t in timestamps]
timestamps = list(map(int, timestamps))  # Convert timestamps to integers

enemy_timestamps = speed_data["enemy_speeds"]
enemy_x_speeds = [speed_data["enemy_speeds"][t][0] for t in enemy_timestamps]
enemy_y_speeds = [speed_data["enemy_speeds"][t][1] for t in enemy_timestamps]
enemy_z_speeds = [speed_data["enemy_speeds"][t][2] for t in enemy_timestamps]
enemy_timestamps = list(map(int, timestamps))  # Convert timestamps to integers

# Extract control data
control_timestamps = sorted(control_data["control"].keys(), key=int)
control_param_0 = [control_data["control"][str(t)][0] for t in control_timestamps]
control_param_1 = [control_data["control"][str(t)][1] for t in control_timestamps]
control_param_2 = [control_data["control"][str(t)][2] for t in control_timestamps]
control_timestamps = list(map(int, control_timestamps))  # Convert timestamps to integers

# Plotting with dual y-axes and multiple control curves
fig, ax1 = plt.subplots(figsize=(10, 6))

# Speed data on the left y-axis
ax1.plot(timestamps, x_speeds, label="X Speed", marker=',', color='blue')
ax1.plot(timestamps, y_speeds, label="Y Speed", marker=',', color='green')
ax1.plot(timestamps, z_speeds, label="Z Speed", marker=',', color='red')

ax1.plot(timestamps, enemy_x_speeds, label="Enemy X Speed", marker=',', color='blue', linestyle='-.')
ax1.plot(timestamps, enemy_y_speeds, label="Enemy Y Speed", marker=',', color='green', linestyle='-.')
ax1.plot(timestamps, enemy_z_speeds, label="Enemy Z Speed", marker=',', color='red', linestyle='-.')

ax1.set_xlabel("Time Steps")
ax1.set_ylabel("Speed", color='black')
ax1.tick_params(axis='y', labelcolor='black')
ax1.legend(loc="upper left")
ax1.grid(True)

# Control data on the right y-axis
ax2 = ax1.twinx()
ax2.plot(control_timestamps, control_param_0, label="Control Param 0", marker='+', color='dimgray', linestyle='--')
ax2.plot(control_timestamps, control_param_1, label="Control Param 1", marker='+', color='orange', linestyle='--')
ax2.plot(control_timestamps, control_param_2, label="Control Param 2", marker='+', color='darkgray', linestyle='--')
ax2.set_ylabel("Control", color='purple')
ax2.tick_params(axis='y', labelcolor='purple')
ax2.legend(loc="upper right")

# Title
plt.title("Real-Time Speeds and Control Commands (All Parameters)")

# 保存图片
plt.savefig("speed_and_control_plot.png", dpi=300, bbox_inches='tight')
print(f"Speed and control trend saved as speed_and_control_plot.png")
plt.show()
